﻿#pragma once
#include <future>
#include <variant>
#include <coroutine>
#include <assert.h>
namespace as
{
	namespace co
	{

		///**
		// * \brief Execute the given awaitable and return a std::promise representing the call.
		// * \note The returned promise will block until the async function finishes or throws.
		// *
		// * This can be used to synchronously wait for the result of a coroutine.
		// */
		//template<typename T, typename Awaitable>
		//std::future<T> as_promise(Awaitable&& awaiter)
		//	requires(!std::is_void_v<T>)
		//{
		//	std::promise<T> promise;
		//	auto res = promise.get_future();
		//	[](std::decay_t<Awaitable> awaiter, std::promise<T> promise) -> TaskPromise<> {
		//		try {
		//			promise.set_value(co_await std::move(awaiter));
		//		}
		//		catch (...) { promise.set_exception(std::current_exception()); }
		//		}(std::forward<decltype(awaiter)>(awaiter), std::move(promise));
		//	return res;
		//}

		///**
		// * \brief Execute the given awaitable and return a std::promise representing the call.
		// * \note The returned promise will block until the async function finishes or throws.
		// *
		// * This can be used to synchronously wait for the result of a coroutine.
		// */
		//template<typename T, typename Awaitable>
		//std::future<void> as_promise(Awaitable&& awaiter)
		//	requires(std::is_void_v<T>)
		//{
		//	std::promise<void> promise;
		//	auto res = promise.get_future();
		//	[](std::decay_t<Awaitable> awaiter, std::promise<void> promise) -> TaskPromise<> {
		//		try {
		//			co_await std::move(awaiter);
		//			promise.set_value();
		//		}
		//		catch (...) { promise.set_exception(std::current_exception()); }
		//		}(std::forward<decltype(awaiter)>(awaiter), std::move(promise));
		//	return res;
		//}

		/**
		 * \brief Execute the given awaitable and return a std::promise representing the call.
		 * \note The returned promise will block until the async function finishes or throws.
		 * \note This function tries to autodetect the return type of the co_await expression.
		 * However the C++20 coroutines have lots of extension points which makes this hard (impossible?)
		 * to do in all cases. If the detection fails you can use the above functions and explicitly specify the return type.
		 *
		 * This can be used to synchronously wait for the result of a coroutine.
		 */
		 //auto as_promise(auto&& awaiter) {
		 //	return as_promise<typename detail::await_return_type<std::remove_cvref_t<decltype(awaiter)>>::type>(
		 //		std::forward<decltype(awaiter)>(awaiter));
		 //}

		template<class T, class TaskPromise>
		class TaskPromiseBase
		{
		public:
			auto get_return_object() noexcept {
				return std::coroutine_handle<TaskPromise>::from_promise(*static_cast<TaskPromise*>(this));
			}

			std::suspend_always initial_suspend() { return {}; }
			auto final_suspend() noexcept {
				struct awaiter {
					constexpr bool await_ready() noexcept { return false; }
					auto await_suspend(std::coroutine_handle<TaskPromise> hndl) noexcept {
						assert(hndl);
						assert(hndl.promise().handle_);
						return hndl.promise().handle_;
					}
					constexpr void await_resume() noexcept {}
				};
				return awaiter{};
			}
		private:
			std::coroutine_handle<> handle_{};
			std::variant<T, int> var_;
		};

		template<class T>
		class TaskPromise : public TaskPromiseBase<T, TaskPromise<T> >
		{
		public:
			void return_value(T&& val)
			{
				this->var_ = std::move(val);
			}

			T get()
			{
				return this->var_;
			}
		};

		struct returned {};
		template<>
		class TaskPromise<void> : public TaskPromiseBase<returned, TaskPromise<void> >
		{
		public:
			void return_void() {};
			void get() {};
		};


		template<class T>
		class CoTask
		{
		public:
			using promise_type = TaskPromise<T>;
			using promise_handle = std::coroutine_handle<promise_type>;

			auto operator co_await() noexcept {
				struct awaiter {
					constexpr explicit awaiter(promise_handle coro) : m_coro(coro) {}
					constexpr bool await_ready() noexcept { return false; }
					auto await_suspend(std::coroutine_handle<void> hndl) noexcept {
						assert(this->m_coro);
						assert(hndl);
						m_coro.promise().handle = hndl;
						return m_coro;
					}
					T await_resume() {
						assert(this->m_coro);
						return this->m_coro.promise().get();
					}

				private:
					promise_handle m_coro;
				};
				assert(this->m_coro);
				return awaiter{ m_coro };
			}

		private:
			promise_handle m_coro;
		};

	}
}
